﻿using System;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;

namespace RayDen.Library.Components.Surface.Microfacet
{
    public class Anisotropic : MicrofacetDistribution
    {
        private float ex, ey;

        public Anisotropic(float x, float y)
        {
            ex = x; ey = y;
            if (ex > 1000f || float.IsNaN(ex)) ex = 1000f;
            if (ey > 1000f || float.IsNaN(ey)) ey = 1000f;
        }

        private void sampleFirstQuadrant(float u1, float u2, out float phi, out float costheta)
        {
            if (ex.NearEqual(ey))
                phi = MathLab.M_PI * u1 * 0.5f;
            else
                phi = (float)Math.Atan(MathLab.Sqrt((ex + 1) / (ey + 1)) *
                                       Math.Tan(MathLab.M_PI * u1 * 0.5f));
            float cosphi = MathLab.Cos(phi), sinphi = MathLab.Sin(phi);
            costheta = (float)MathLab.Pow(u2, 1f / (ex * cosphi * cosphi + ey * sinphi * sinphi + 1));
        }

        public override float D(ref Vector wh)
        {
            float costhetah = Math.Abs(BrdfBase.CosTheta(ref wh));
            float e = (ex * wh.x * wh.x + ey * wh.y * wh.y) /
                      (1f - costhetah * costhetah);
            return MathLab.Sqrt((ex + 2) * (ey + 2)) * MathLab.INVTWOPI * (float)MathLab.Pow(costhetah, e);
        }

        public override void Sample_f(ref Vector wo, out Vector wi, float u1, float u2, out float pdf)
        {
            float phi, costheta;
            if (u1 < .25f)
            {
                sampleFirstQuadrant(4f * u1, u2, out phi, out costheta);
            }
            else if (u1 < .5f)
            {
                u1 = 4f * (.5f - u1);
                sampleFirstQuadrant(u1, u2, out phi, out costheta);
                phi = MathLab.M_PI - phi;
            }
            else if (u1 < .75f)
            {
                u1 = 4f * (u1 - .5f);
                sampleFirstQuadrant(u1, u2, out phi, out costheta);
                phi += MathLab.M_PI;
            }
            else
            {
                u1 = 4f * (1f - u1);
                sampleFirstQuadrant(u1, u2, out phi, out costheta);
                phi = 2f * MathLab.M_PI - phi;
            }
            float sintheta = MathLab.Sqrt(Math.Max(0f, 1f - costheta * costheta));
            Vector wh = MathLab.SphericalDirection(sintheta, costheta, phi);
            if (!BrdfBase.SameHemisphere(ref wo, ref wh)) wh = -wh;

            // Compute incident direction by reflecting about $\wh$
            wi = -wo + 2f * Vector.Dot(ref wo, ref wh) * wh;

            // Compute PDF for $\wi$ from anisotropic distribution
            float costhetah = BrdfBase.AbsCosTheta(ref wh);
            float ds = 1f - costhetah * costhetah;
            float anisotropicPdf = 0f;
            if (ds > 0f && Vector.Dot(ref wo, ref wh) > 0f)
            {
                float e = (ex * wh.x * wh.x + ey * wh.y * wh.y) / ds;
                float d = MathLab.Sqrt((ex + 1f) * (ey + 1f)) * MathLab.INVTWOPI *
                          (float)MathLab.Pow(costhetah, e);
                anisotropicPdf = d / (4f * Vector.Dot(ref wo, ref wh));
            }
            pdf = anisotropicPdf;
        }

        public override float Pdf(ref Vector wo, ref Vector wi)
        {
            Vector wh = Vector.Normalize(wo + wi);
            // Compute PDF for $\wi$ from anisotropic distribution
            float costhetah = BrdfBase.AbsCosTheta(ref wh);
            float ds = 1f - costhetah * costhetah;
            float anisotropicPdf = 0f;
            if (ds > 0f && Vector.Dot(ref wo, ref  wh) > 0f)
            {
                float e = (ex * wh.x * wh.x + ey * wh.y * wh.y) / ds;
                float d = MathLab.Sqrt((ex + 1f) * (ey + 1f)) * MathLab.INVTWOPI * (float)MathLab.Pow(costhetah, e);
                anisotropicPdf = d / (4f * Vector.Dot(ref wo, ref wh));
            }
            return anisotropicPdf;
        }
    }
}